import random
import numpy as np
import math
import json
import xml.etree.ElementTree as ET
import os


# mjcf示例
mjcf_template = """
<mujoco model="template">
  <option gravity="0 0 -9.8" />
  <statistic center="0 0 0.1" extent="0.8"/>
  <visual>
    <headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0"/>
    <rgba haze="0.15 0.25 0.35 1"/>
    <global azimuth="-130" elevation="-20"/>
    <map znear="0.1" zfar="10000"/>
  </visual>
  <asset>
    <texture type="skybox" builtin="gradient" rgb1="0.3 0.5 0.7" rgb2="0 0 0" width="512" height="3072"/>
    <texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3"
      markrgb="0.8 0.8 0.8" width="300" height="300"/>
    <material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance="0.2"/>
    <material name="SlateGrey" rgba="0.4392 0.5020 0.5647 1" reflectance="0.2" />
    <material name="translucent_black" rgba="0 0 0 0.5"/>
    <material name="translucent_red" rgba="1 0 0 0.5"/>
    <material name="translucent_green" rgba="0 1 0 0.5"/>
    <material name="translucent_blue" rgba="0 0 1 0.5"/>
  </asset>
  <worldbody>
    <light pos="0 0 1.5" dir="0 0 -1" directional="true"/>
    <geom name="floor" size="10 10 0.05" type="plane" material="groundplane" friction="1.0 0.02 0.01"/>
  </worldbody>
</mujoco>
"""

def _z_rotation_quat(angle):
    """生成绕Z轴旋转的四元数"""
    # 四元数公式: (cos(θ/2), 0, 0, sin(θ/2))
    half_angle = angle / 2
    w = math.cos(half_angle)
    z = math.sin(half_angle)
    return [w, 0, 0, z]  # [w, x, y, z]

class StairwellScene:
    def __init__(
        self, 
        scnens_size: list = [10, 10], 
        collision_num: int = 10,
        scene_params: dict = {},
        collision_size_limit: dict = {
            "box":{
                "x": [0.1, 0.75],
                "y": [0.1, 1]
            },
            "cylinder":{
                "r": [0.1, 0.75]
            }
        }
    ):
        self.scnens_size = scnens_size
        self.collision_num = collision_num
        self.collision_size_limit = collision_size_limit
        self.scene_params = scene_params
        self.collision_config = [
            {
                "type": "box",
                "size": [1, 1, 1],
                "pos": [0, 0, 0],
                "quat": [1, 0, 0, 0]
            },
            {
                "type": "cylinder",
                "size": [1, 1],
                "pos": [0, 0, 0],
                "quat": [1, 0, 0, 0]
            }
        ]

    def gen(self):
        self.collision_config = []
        center_aisle_with = random.uniform(3, 6)
        side_aisle_with = random.uniform(3, 6)
        self.gen_wall(center_aisle_with, side_aisle_with)
        self.scene_params["size"] = self.scnens_size
        self.scene_params["wall"] = {
            "center_aisle_with": center_aisle_with,
            "side_aisle_with": side_aisle_with
        }
        
        for _ in range(self.collision_num):
            obj_type = random.choice(["box", "cylinder"])
            if obj_type == "box":
                size = [
                    random.uniform(
                        self.collision_size_limit["box"]["x"][0], 
                        self.collision_size_limit["box"]["x"][1]
                    ),
                    random.uniform(
                        self.collision_size_limit["box"]["y"][0], 
                        self.collision_size_limit["box"]["y"][1]
                    ),
                    1.5
                ]
                z_pos = size[2]
            else:  # cylinder
                size = [
                    random.uniform(
                        self.collision_size_limit["cylinder"]["r"][0], 
                        self.collision_size_limit["cylinder"]["r"][1]
                    ),
                    1.5
                ]
                z_pos = size[1]
            pos = self.get_random_pos_on_wall(center_aisle_with, side_aisle_with, z_pos)
            
            # 生成绕Z轴的随机旋转（0-360度）
            angle = random.uniform(0, 2 * math.pi)  # 随机角度（弧度）
            quat = _z_rotation_quat(angle)  # 转换为四元数
            # 添加到配置列表
            self.collision_config.append(
                {
                    "type": obj_type,
                    "size": size,
                    "pos": pos,
                    "quat": quat
                }
            )
        return self.collision_config, self.scene_params

    def get_random_pos_on_wall(self, center_aisle_with, side_aisle_with, z_pos):
        pos = random.choice(
            [
                [
                    random.uniform(
                        -self.scnens_size[0] / 2, self.scnens_size[0] / 2
                    ),
                    random.uniform(
                        center_aisle_with / 2 - 0.3, center_aisle_with / 2
                    ),
                    z_pos
                ],
                [
                    random.uniform(
                        -self.scnens_size[0] / 2, -side_aisle_with / 2
                    ),
                    random.uniform(
                        -center_aisle_with / 2, -center_aisle_with / 2 + 0.3
                    ),
                    z_pos
                ],
                [
                    random.uniform(
                        side_aisle_with / 2, self.scnens_size[0] / 2
                    ),
                    random.uniform(
                        -center_aisle_with / 2, -center_aisle_with / 2 + 0.3
                    ),
                    z_pos
                ],
                [
                    random.uniform(
                        -side_aisle_with / 2, -side_aisle_with / 2 + 0.3
                    ),
                    random.uniform(
                        -self.scnens_size[1] / 2, -center_aisle_with / 2
                    ),
                    z_pos
                ],
                [
                    random.uniform(
                        side_aisle_with / 2 - 0.3, side_aisle_with / 2
                    ),
                    random.uniform(
                        -self.scnens_size[1] / 2, -center_aisle_with / 2
                    ),
                    z_pos
                ]
            ]
        )
        return pos
        
    # 生成过道的墙体
    def gen_wall(self, center_aisle_with, side_aisle_with):
        self.collision_config.append(
            {
                "type": "box",
                "size": [self.scnens_size[0] / 2, 0.05, 1.5],
                "pos": [0, center_aisle_with / 2, 1.5],
                "quat": [1, 0, 0, 0]
            }
        )
        length = self.scnens_size[0] / 2 - side_aisle_with / 2
        self.collision_config.append(
            {
                "type": "box",
                "size": [length / 2 + 0.025, 0.05, 1.5],
                "pos": [-side_aisle_with / 2 - length / 2 + 0.025, -center_aisle_with / 2, 1.5],
                "quat": [1, 0, 0, 0]
            }
        )
        self.collision_config.append(
            {
                "type": "box",
                "size": [length / 2 + 0.025, 0.05, 1.5],
                "pos": [side_aisle_with / 2 + length / 2 - 0.025, -center_aisle_with / 2, 1.5],
                "quat": [1, 0, 0, 0]
            }
        )
        length = self.scnens_size[1] / 2 - center_aisle_with / 2
        self.collision_config.append(
            {
                "type": "box",
                "size": [0.05, length / 2 + 0.025, 1.5],
                "pos": [-side_aisle_with / 2, -center_aisle_with / 2 - length / 2 + 0.025, 1.5],
                "quat": [1, 0, 0, 0]
            }
        )
        self.collision_config.append(
            {
                "type": "box",
                "size": [0.05, length / 2 + 0.025, 1.5],
                "pos": [side_aisle_with / 2, -center_aisle_with / 2 - length / 2 + 0.025, 1.5],
                "quat": [1, 0, 0, 0]
            }
        )
        
        

class RandomScene:
    def __init__(
        self, 
        scnens_size: list = [10, 10], 
        collision_num: int = 10,
        collision_size_limit: dict = {
            "box":{
                "x": [0.1, 1],
                "y": [0.1, 1]
            },
            "cylinder":{
                "r": [0.1, 1]
            },
            "wall":{
                "length": [1, 2]
            }
        }
    ):
        self.scnens_size = scnens_size
        self.collision_num = collision_num
        self.collision_size_limit = collision_size_limit
        self.collision_config = [
            {
                "type": "box",
                "size": [1, 1, 1],
                "pos": [0, 0, 0],
                "quat": [1, 0, 0, 0]
            },
            {
                "type": "cylinder",
                "size": [1, 1],
                "pos": [0, 0, 0],
                "quat": [1, 0, 0, 0]
            }
        ]
    
    def gen(self):
        self.collision_config = []
        # 地图边界（中心为原点）
        x_min, x_max = -self.scnens_size[0]/2, self.scnens_size[0]/2
        y_min, y_max = -self.scnens_size[1]/2, self.scnens_size[1]/2
        
        for _ in range(self.collision_num):
            # 随机选择物体类型 (box 或 cylinder)
            obj_type = random.choice(["box", "cylinder", "wall"])
            
            if obj_type == "box":
                # 随机生成box尺寸 (半长, 半宽, 半高) 范围: [0.1, 1.0]
                size = [
                    random.uniform(
                        self.collision_size_limit["box"]["x"][0], 
                        self.collision_size_limit["box"]["x"][1]
                    ),
                    random.uniform(
                        self.collision_size_limit["box"]["y"][0], 
                        self.collision_size_limit["box"]["y"][1]
                    ),
                    1.5
                ]
                # Z位置设为半高，使物体底部在地面
                z_pos = size[2]
            elif obj_type == "cylinder":  # cylinder
                # 随机生成cylinder尺寸 (半径, 半高) 范围: [0.1, 1.0]
                size = [
                    random.uniform(
                        self.collision_size_limit["cylinder"]["r"][0], 
                        self.collision_size_limit["cylinder"]["r"][1]
                    ),
                    1.5
                ]
                # Z位置设为半高
                z_pos = size[1]
            else:  # wall
                size = [
                    random.uniform(
                        self.collision_size_limit["wall"]["length"][0], 
                        self.collision_size_limit["wall"]["length"][1]
                    ),
                    0.05,
                    1.5
                ]
                z_pos = size[2]
            
            # 在地图范围内随机生成XY位置
            x_pos = random.uniform(x_min, x_max)
            y_pos = random.uniform(y_min, y_max)
            
            # 生成绕Z轴的随机旋转（0-360度）
            angle = random.uniform(0, 2 * math.pi)  # 随机角度（弧度）
            quat = _z_rotation_quat(angle)  # 转换为四元数
            
            if obj_type == "wall":
                obj_type = "box"
            # 添加到配置列表
            self.collision_config.append(
                {
                    "type": obj_type,
                    "size": size,
                    "pos": [x_pos, y_pos, z_pos],
                    "quat": quat
                }
            )
        
        return self.collision_config
        



class GenMJCFScenes:
    def __init__(
        self, 
        mode:str = "random",
        collision_num: int = 10,
        gen_scene_num: int = 1,
        scene_save_dir: str = "dataset_tools/sources/mjcf_scenes"
    ):
        self.mode = mode
        self.collision_num = collision_num
        self.scene_save_dir = scene_save_dir
        self.gen_scene_num = gen_scene_num
        
    # 生成场景
    def gen(self):
        for _ in range(self.gen_scene_num):
            scene_params = {
                "mode": self.mode
            }
            self.mjcf = self.get_mjcf_template()
            self.scnens_size = self.get_scene_size()

            self.scene_mark = self.get_sence_num_mark()
            self.mjcf.set("model", f"scene_{self.scene_mark}")
            
            collision_config, scene_params = self.get_collision_config(scene_params)
            worldbody = self.mjcf.find('worldbody')
            worldbody = self.gen_wall(worldbody, self.scnens_size)
            for collision in collision_config:
                ET.SubElement(worldbody, "geom", {
                    "type": collision["type"],
                    "size": self.list2str(collision["size"]),
                    "pos": self.list2str(collision["pos"]),
                    "quat": self.list2str(collision["quat"]),
                    "material": "SlateGrey",
                    "group": "1"
                })
            self.save_scene(scene_params)

    def get_collision_config(self, scene_params):
        if self.mode == "random":
            random_scene = RandomScene(
                scnens_size=self.scnens_size,
                collision_num=self.collision_num
            )
            return random_scene.gen(), scene_params
        elif self.mode == "stairwell":
            random_scene = StairwellScene(
                scnens_size=self.scnens_size,
                collision_num=self.collision_num,
                scene_params=scene_params
            )
            return random_scene.gen()
        else:
            raise KeyError(f"Not support mode {self.mode}")
        
    def get_mjcf_template(self):
        mjcf = ET.fromstring(mjcf_template)
        return mjcf

    # 获取scene标号，scene_0, scene_1, scene_2...
    def get_sence_num_mark(self):
        scenes_dir_list = [
            dir 
            for dir in os.listdir(self.scene_save_dir) 
            if os.path.isdir(os.path.join(self.scene_save_dir, dir)) and "scene_" in dir
        ]
        scenes_mark = []
        for scene_dir in scenes_dir_list:
            scene_num = int(scene_dir.split("_")[1])
            scenes_mark.append(scene_num)
        return max(scenes_mark) + 1 if len(scenes_mark) > 0 else 0
    
    # 保存xml
    def save_scene(self, scene_params):
        scene_save_dir = os.path.join(self.scene_save_dir, f"scene_{self.scene_mark}")
        os.makedirs(scene_save_dir, exist_ok=True)
        save_mjcf = os.path.join(scene_save_dir, f"scene_{self.scene_mark}.xml")
        with open(save_mjcf, "w", encoding="utf-8") as f:
            f.write('<?xml version="1.0" encoding="utf-8"?>\n')
            ET.indent(self.mjcf, space="    ", level=0)  # Python 3.9+ 才有的缩进功能
            f.write(ET.tostring(self.mjcf, encoding="unicode"))
        scene_params_path = os.path.join(scene_save_dir, f"scene_{self.scene_mark}.json")
        with open(scene_params_path, "w", encoding="utf-8") as f:
            json.dump(scene_params, f, indent=4, ensure_ascii=False)
            
    def list2str(self, lst):
        s = ""
        for i, val in enumerate(lst):
            s += f"{val}"
            if i != len(lst) - 1:
                s += " "
        return s
    # 根据floor获取地图尺寸
    def get_scene_size(self):
        floor_geom = self.mjcf.find(".//worldbody/geom[@name='floor']")
        if floor_geom is not None:
            size = [float(val) * 2 for val in floor_geom.get("size").split(" ")]
            return size
        return [20, 20, 0.1]
    
    # 生成地图四周的围墙
    def gen_wall(self, worldbody, scnens_size):
        ET.SubElement(worldbody, "geom", 
            {
                "name": "wall1",
                "type": "box",
                "size": f"0.05 {scnens_size[1] / 2} 1.5",
                "pos": f"{scnens_size[0] / 2} 0 1.5",
                "material": "SlateGrey",
                "group": "1"
            }
        )
        ET.SubElement(worldbody, "geom", 
            {
                "name": "wall2",
                "type": "box",
                "size": f"0.05 {scnens_size[1] / 2} 1.5",
                "pos": f"-{scnens_size[0] / 2} 0 1.5",
                "material": "SlateGrey",
                "group": "1"
            }
        )
        ET.SubElement(worldbody, "geom", 
            {
                "name": "wall3",
                "type": "box",
                "size": f"{scnens_size[0] / 2} 0.05 1.5",
                "pos": f"0 {scnens_size[1] / 2} 1.5",
                "material": "SlateGrey",
                "group": "1"
            }
        )
        ET.SubElement(worldbody, "geom", 
            {
                "name": "wall4",
                "type": "box",
                "size": f"{scnens_size[0] / 2} 0.05 1.5",
                "pos": f"0 -{scnens_size[1] / 2} 1.5",
                "material": "SlateGrey",
                "group": "1"
            }
        )
        return worldbody
            


if __name__ == "__main__":
    genMJCFScenes = GenMJCFScenes(
        # mode="stairwell",
        mode="random",
        collision_num=15,
        gen_scene_num=15
    )
    genMJCFScenes.gen()